#include "KalmanPredictor.h"
using namespace robowalker;

KalmanPredictor::KalmanPredictor(const double Dt)
{
    dt = Dt;
    double ddt = 0.5 * (dt * dt);
    angleX_measurement_queue = deque<double>();
    angleY_measurement_queue = deque<double>();

    B = Mat_<double>(4, 2)
        << ddt,
    0,
    0, ddt,
    dt, 0,
    0, dt;

    initKF();
}

void KalmanPredictor::initKF()
{
    double Init_Q = 1e-5, Init_R = 1e-1; //初始参数，后面需要调参

    // State : X : 4 x 1
    // measure: Z : 4 x 1
    KF.init(4, 4, 2, CV_32F);
    KF.transitionMatrix = Mat_<double>(4, 4)
                          << 1,
    0, dt, 0,
    0, 1, 0, dt,
    0, 0, 1, 0,
    0, 0, 0, 1;
    KF.transitionMatrix.convertTo(KF.transitionMatrix, CV_32F);
    // KF.measurementMatrix = Mat_<double>(4, 2)
    //         <<  1, 0, 0, 0,
    //             0, 1, 0, 0,
    //             0, 0, 1, 0,
    //             0, 0, 0, 1;
    // KF.measurementMatrix.convertTo(KF.measurementMatrix, CV_32F);
    // KF.processNoiseCov = Mat_<double>(4, 4)
    //         <<  Init_Q, 0, 0, 0,  // Q: process noise covariance : #TODO adjust
    //         0, Init_Q, 0, 0,
    //         0, 0, Init_Q, 0,
    //         0, 0, 0, Init_Q;
    // KF.measurementNoiseCov = Mat_<double>(2, 2)
    //         << Init_R, 0,  // R: measurement noise covariance: #TODO adjust
    //         0, Init_R;

    setIdentity(KF.measurementMatrix);                      //测量矩阵H
    setIdentity(KF.processNoiseCov, Scalar::all(1e-5));     //系统噪声方差矩阵Q
    setIdentity(KF.measurementNoiseCov, Scalar::all(1e-6)); //测量噪声方差矩阵R
    setIdentity(KF.errorCovPost, Scalar::all(1));           //后验错误估计协方差矩阵P
    KF.controlMatrix = B;
    KF.controlMatrix.convertTo(KF.controlMatrix, CV_32F);

    randn(KF.statePost, Scalar::all(0), Scalar::all(0.1)); // Initialize
}

void KalmanPredictor::reset()
{
    initKF();
}

Mat KalmanPredictor::update(Mat measurement, Mat U)
{
    // printf("Kalman update\n");
    // Mat control = B * U;
    // control.convertTo(control, CV_32F);
    // U.convertTo(U, CV_32F);
    auto X_predict = KF.predict();
    // assert(X_predict.size.dims() == 2);
    // printf("Kalman correct\n");
    X_predict = KF.correct(measurement);
    return X_predict; // 4 x 1
}

/*
 * @brief: add measurement angle to queue
 */
void KalmanPredictor::addMeasurementToQueue(const double angleX_measurement, const double angleY_measurement)
{
    if (angleX_measurement_queue.size() >= 20)
        angleX_measurement_queue.pop_front();
    angleX_measurement_queue.push_back(angleX_measurement);

    if (angleY_measurement_queue.size() >= 20)
        angleY_measurement_queue.pop_front();
    angleY_measurement_queue.push_back(angleY_measurement);
}

void KalmanPredictor::clearMeasurementQueue()
{
    angleX_measurement_queue.clear();
    angleY_measurement_queue.clear();
}

pair<double, double> KalmanPredictor::velocityPredict() const
{
    int n = angleX_measurement_queue.size();
    const int n_thresh = 15; 
    if (n > n_thresh)
    {
        const double time_step = 5.0; // 5 ms

        Mat_<double> A_yaw = Mat_<double>::ones(n, 2);
        Mat_<double> A_pitch = Mat_<double>::ones(n, 2);
        Mat_<double> b_yaw = Mat_<double>::zeros(n, 1);
        Mat_<double> b_pitch = Mat_<double>::zeros(n, 1);

        for (int i = 0; i < n; i++)
        {
            b_yaw(i) = angleX_measurement_queue[i];
            b_pitch(i) = angleY_measurement_queue[i];

            A_yaw(i, 0) = A_pitch(i, 0) = time_step * i;
        }

        Mat_<double> X_yaw, X_pitch;

        solve(A_yaw.t() * A_yaw, A_yaw.t() * b_yaw, X_yaw); // least square solve
        solve(A_pitch.t() * A_pitch, A_pitch.t() * b_pitch, X_pitch);

        double v_yaw = X_yaw(0, 0);
        double v_pitch = X_pitch(0, 0);
        return {v_yaw, v_pitch}; // metric: degree per ms
    }
    else
    {
        return {0.0, 0.0};
    }
}

void robowalker::thread_kalmanPredictor(const Mat *ptr_angleMeasurement, Mat *ptr_predict_angle, const bool *ptr_reset,
                                        const double *dist, car_info *car_data, int *detectState)
{
    printf("[Thread] KalmanPredictor thread starts. \n");
    const double kalmanDt = 5; //
    KalmanPredictor KP(kalmanDt);
    const char *manifoldSerialPath = "/dev/ttyTHS2";
    const int baudrate = 115200;
    const double bullet_speed = 30;       // metric: mm/ms;
    const double feed_bullet_delay = 100; // metric: ms;
    const double v_range = 12.0 / 1000;   // degree per ms
    CommunicationManager CM(manifoldSerialPath, baudrate);

    while (1)
    {
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
        if (*ptr_reset)
            KP.reset();

        CM.manifoldReceiveOnce(car_data);

        // 需要减去云台自己的运动角速度
        const double k = 30.0 / 37.0;
        float self_vy = float(car_data->vy) * k / 10000.0 ;
        float self_vx = float(car_data->vx) * k / 10000.0 ; 

        auto angleX = (*ptr_angleMeasurement).at<float>(0);
        auto angleY = (*ptr_angleMeasurement).at<float>(1);

        KP.addMeasurementToQueue(angleX, angleY);
        double v_x = 0.0;
        double v_y = 0.0;
        
        if( (*detectState) == 1)
        {
            auto velocity = KP.velocityPredict();
            v_x = velocity.first;
            v_y = velocity.second;
        }
        else
        {
            KP.clearMeasurementQueue();
        }
        // printf("vx: %lf, vy: %lf\n", v_x, v_y);

        Mat measure = (Mat_<double>(4, 1) << angleX, angleY, v_x, v_y);
        measure.convertTo(measure, CV_32F);

        // auto U = (Mat_<double>(2,1)<< v_x, v_y);

        auto result = KP.update(measure);

        auto predict_angleX = result.at<float>(0);
        auto predict_angleY = result.at<float>(1);

        auto predict_vx = result.at<float>(2) ;//+ self_vx;
        auto predict_vy = result.at<float>(3) ;//+ self_vy;
        // printf("predict: %lf, self: %lf\n", predict_vx, self_vx);

        double new_predict_angleX, new_predict_angleY; 
        
        if( (*detectState) == 1)
        {
            new_predict_angleX = predict_angleX + predict_vx * ((*dist) / bullet_speed + feed_bullet_delay);
            new_predict_angleY = predict_angleY + predict_vy * ((*dist) / bullet_speed + feed_bullet_delay);
        }
        else 
            {
                 new_predict_angleX = 0.0;
                 new_predict_angleY = 0.0;
            }
        (*ptr_predict_angle).at<float>(0) = new_predict_angleX;
        (*ptr_predict_angle).at<float>(1) = new_predict_angleY;

        double c = 10.0;

        // if(angleX > c || angleX < -c) printf("[Error] Wrong!");
        // if(angleY > c || angleY < -c) printf("[Error] Wrong!");
        // if(predict_angleX > c || predict_angleX < -c) printf("[Error] Wrong!");
        // if(predict_angleY > c || predict_angleY < -c) printf("[Error] Wrong!");

        // printf("[KF thread]\n");
        // printf("angle measurement X: %f , angle measurement Y: %f \n", angleX, angleY);
        // printf("angle prediction X: %f , angle prediction Y: %f \n\n", predict_angleX, predict_angleY);

        CM.manifoldSendOnce(assembleVisionInfoStruct(angleX, angleY, new_predict_angleX, new_predict_angleY, *dist, predict_vx, predict_vy));
    }
}